# -*- coding: utf-8 -*-
"""
A cultural systems model, accompanying the manuscript
"A Systems Approach to Cultural Evolution".

@author: Fredrik Jansson
"""

import numpy as np
from random import random, choice, randrange

class Model:
    """
    Simulates agents adopting cultural traits drawn from a trait pool/universe.
    Traits can be compatible to various degrees. Agents aquire traits by
    copying from others or by randomly drawing traits from the pool. A trait
    is more likely to be copied if it is compatible with other traits that
    the agent possesses.
    """
    
    def __init__(self,pool,population_size,generation_length,p_invent,p_copy=1,k=10):
        """
        Parameters
        ----------
        pool : Pool
            The trait universe
        population_size : int
            Number of agents
        generation_length : int
            Expected number of interactions before an agent is replaced
        p_invent : double
            Probability for agent to sample a trait from the pool
        p_copy : double, optional
            Scales the probability of copying traits from another agent
        k : double, optional
            Increases the slope of the s-shape in the logitistic function
            determining the probability to copy
        """
        self.pool = pool
        self.population_size = population_size
        self.pool_size = self.pool.size
        self.population = Population(population_size,self.pool_size)
        self.generation_length = generation_length
        self.p_invent = p_invent
        self.p_copy = p_copy
        self.trait_filter = True
        self.k = k
        self.time_step = int()
    
    def simulate(self,iterations):
        """
        Runs a simulation for iterations number of time steps. The simulation
        runs from the present state, i.e. if the last run was until time step
        t, then the first time step here is t+1. In each iteration, each agent
        randomly selects one model agent for an interaction, potentially
        invents a new trait and is replaced by a naïve individual with an
        inverse probability to the generation length.
        
        Parameters
        ----------
        iterations : int
            Number of time steps
        """
        for i in range(iterations):
            for k, agent in enumerate(self.population.agents):
                self.interact(agent,self.population.random_agent())
                self.invent(agent)
                if random() < 1 / self.generation_length:
                    self.population.replace_agent(k) # death and birth
            self.time_step += 1
    
    def repeat(self,iterations,runs,measures):
        """
        Calls the function simulate repeatedly.
        
        Parameters
        ----------
        iterations : int
            Number of time steps (passed to simulate)
        runs : int
            Number of consecutive calls to simulate
        measures : [function]
            Measures to compute after each call to simulate
        
        Returns
        -------
        (list, numpy.array)
            A list of time steps and an array of measure computed by the
            measures functions
        """
        results = []
        time_steps = np.zeros(runs)
        for i in range(runs):
            self.simulate(iterations)
            results.extend([measure(self) for measure in measures])
            time_steps[i] = self.time_step
        return (time_steps,np.array(results))

    def interact(self,observer,sender):
        """
        An agent observes a random trait in another agent an copies it with
        probability proportional to how compatible it is with all the other
        traits in the observing agent's current repertoire.
        
        Parameters
        ----------
        observer : Agent
            The observing agent
        sender : Agent
            The agent to observe and potentially copy from
        """
        if sender.has_culture():
            observed_trait = sender.random_trait()
            if not observer.has_trait(observed_trait):
                if observer.has_culture():
                    s = np.nanmean(self.pool.compatibility[observer.has_traits,observed_trait]) # compatibility score
                    if np.isnan(s): # compatibility is undefined
                        s = 0
                else:
                    s = 0
                p = self.p_copy / (1 + np.exp(-self.k*s)) # probability of copying
                if random() < p:
                    observer.add_trait(observed_trait)

    def invent(self,agent):
        """
        Sample a trait from the pool with probability p_invent.
        
        Parameters
        ----------
        agent : Agent
            The agent that will potentially invent a trait added to their
            repertoire
        """
        if random() < self.p_invent:
            innovation = self.pool.random_trait()
            if not agent.has_trait(innovation):
                agent.add_trait(innovation)


class Population:
    """
    A population of agents with cultural traits.
    """
    
    def __init__(self,population_size,pool_size):
        """
        Parameters
        ----------
        population_size : int
            Number of agents in the population
        pool_size : int
            Size of the trait universe, used for efficient managing of agents'
            repertoires
        """
        self.population_size = population_size
        self.pool_size = pool_size
        self.agents = [Agent(pool_size) for _ in range(population_size)]
       
    def get_agents(self):
        """
        Returns
        -------
        [Agent]
            All agents in the population
        """
        return self.agents
    
    def random_agent(self):
        """
        Returns
        -------
        Agent
            A random agent from the population
        """
        return self.agents[randrange(self.population_size)]
    
    def replace_agent(self,k):
        """
        Parameters
        ----------
        k : int
            Index of agent to replace
        """
        self.agents[k] = Agent(self.pool_size)
    
    def has_traits(self):
        """
        Returns
        -------
        [bool]
            A list indicating whether each agent has a cultural trait
        """
        return np.array([agent.has_traits for agent in self.agents])
    

class Pool:
    """
    A cultural trait pool/universe. Traits are connected, specified by an
    adjacency matrix.
    """
    
    def __init__(self,adj_matrix):
        """
        Parameters
        ----------
        adj_matrix : [[double]]
            Adjacency matrix specifying compatibilities between each pair of
            cultural traits
        """
        self.compatibility = np.array(adj_matrix)
        self.size = self.compatibility.shape[0]
        self.compatibility_nan = self.compatibility + np.diag(np.full(self.size,np.nan)) # undefined compatibility to itself
    
    @classmethod
    def random(cls,size,weights,density):
        """
        Create a random trait pool.
        
        Parameters
        ----------
        size : int
            Size of the trait pool
        weights : [double]
            List of compatibilities
        density : [double]
            List of expected proportion of each compatibility listed in weights
        
        Returns
        -------
        Pool
            A trait pool
        """
        compatibility = np.random.choice(np.double(weights),p=density,size=(size,size)) # init trait pool
        np.fill_diagonal(compatibility,1) # traits are compatible to themselves
        compatibility = np.tril(compatibility) + np.tril(compatibility, -1).T # make symmetric
        return cls(compatibility)
    
    def random_trait(self):
        """
        Returns
        -------
        int
            Random trait from the pool
        """
        return randrange(self.size)
    
    def set_compatibilities(self,indices,comps):
        """
        Parameters
        ----------
        indices : [int]
            traits to set mutual compatibilities for
        comps : [[double]]
            matrix of compatibilities with length(indicies) rows and columns
        """
        self.compatibility[indices,indices] = comps
        self.compatibility_nan[indices,indices] = comps + np.diag(np.full(len(comps),np.nan))
    

class Agent:
    """
    An agent with a repertoire of cultural traits.
    """
    
    def __init__(self,pool_size):
        """
        Parameters
        ----------
        pool_size : int
            Size of the trait universe, used for efficient managing of agent's
            repertoire
        """
        self.has_traits = np.zeros(pool_size,dtype=bool) # for quick check and comparison
        self.traits = [] # list has lower complexity for random selection
        self.pool_traits = set(range(pool_size)) # set has lower complexity for set operations
    
    def add_trait(self,nr):
        """
        Parameters
        ----------
        nr : int
            Trait to add to the agent's repertoire
        """
        if not self.has_traits[nr]:
            self.has_traits[nr] = True
            self.traits.append(nr)
            self.pool_traits.remove(nr)
    
    def add_traits(self,nrs):
        """
        Parameters
        ----------
        nrs : [int]
            Traits to add to the agent's repertoire
        """
        [self.add_trait(nr) for nr in nrs]
    
    def remove_trait(self,nr):
        """
        Parameters
        ----------
        nr : int
            Trait to remove from the agent's repertoire
        """
        if self.has_traits[nr]:
            self.has_traits[nr] = False
            self.traits.remove(nr)
            self.pool_traits.add(nr)
    
    def has_culture(self):
        """
        Returns
        -------
        bool
            Whether the agent has at least one trait
        """
        return len(self.traits) > 0
    
    def has_trait(self,nr):
        """
        Parameters
        ----------
        nr : int
            Trait to check
        
        Returns
        -------
        bool
            Whether the agent has trait nr
        """
        return self.has_traits[nr]
    
    def random_trait(self):
        """
        Returns
        -------
        int
            Random trait from the agent's repertoire
        """
        return choice(self.traits)
